import os

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt


def load_df(mode):
    models = ["marank", "sas", "tisas", "bert", "meantime", "modify_meantime"]
    model_list = [model + "_" + mode for model in models]
    root = "experiments/test"
    val_log = "tables/val_log.csv"
    df_paths = [root + "/" + dir_name + "/" + val_log for dir_name in model_list]

    recall_5 = {}
    recall_10 = {}
    ndcg_5 = {}
    ndcg_10 = {}

    for model_name, path in zip(models, df_paths):
        df = pd.read_csv(path)
        print("%s验证记录长度：\t%d" % (model_name.upper(), len(df)))

        df_cache1 = df["Recall@5"].iloc[:50].to_list()
        df_cache2 = df["Recall@10"].iloc[:50].to_list()
        df_cache3 = df["NDCG@5"].iloc[:50].to_list()
        df_cache4 = df["NDCG@10"].iloc[:50].to_list()

        recall_5[model_name] = df_cache1
        recall_10[model_name] = df_cache2
        ndcg_5[model_name] = df_cache3
        ndcg_10[model_name] = df_cache4

    # 创建一个2x2的子图布局
    fig, axs = plt.subplots(2, 2, figsize=(10, 8), dpi=300)

    # 在每个子图对象上绘图
    for model, score in recall_5.items():
        axs[0, 0].plot(range(len(score)), score, label=model)
        axs[0, 0].set_xlabel("Epoch")
        axs[0, 0].set_ylabel("Recall@5")
    axs[0, 0].legend()

    for model, score in recall_10.items():
        axs[0, 1].plot(range(len(score)), score, label=model)
        axs[0, 1].set_xlabel("Epoch")
        axs[0, 1].set_ylabel("Recall@10")
    axs[0, 1].legend()

    for model, score in ndcg_5.items():
        axs[1, 0].plot(range(len(score)), score, label=model)
        axs[1, 0].set_xlabel("Epoch")
        axs[1, 0].set_ylabel("NDCG@5")
    axs[1, 0].legend()

    for model, score in ndcg_10.items():
        axs[1, 1].plot(range(len(score)), score, label=model)
        axs[1, 1].set_xlabel("Epoch")
        axs[1, 1].set_ylabel("NDCG@10")
    axs[1, 1].legend()
    # df_recall_5 = pd.DataFrame(recall_5)
    # df_recall_5.plot()
    # plt.xlabel("Epoch")
    # plt.ylabel("Recall@5")
    # plt.title("")
    for ax in axs.flat:
        ax.set_xlim(0, 50)
        ax.set_ylim(0, 1)
    # 添加整体的xlabel和ylabel
    fig.text(0.5, 0.04, 'Game', ha='center', va='center')
    fig.text(0.04, 0.5, 'Metrics', ha='center', va='center', rotation='vertical')
    plt.savefig(f"{mode}_val_res.png")
    plt.show()
    return


if __name__ == "__main__":
    load_df(mode="game")
    # print(res)

